#linear线性层
#任务：将一张5*5的张量变成1*25的，然后再通过线性层变为1*3的
import torch
import torchvision
from torch import nn
from torch.nn import Linear
from torch.utils.data import DataLoader

dataset = torchvision.datasets.CIFAR10("./dataset",train=False,transform=torchvision.transforms.ToTensor(),download=True)
dataloader = DataLoader(dataset,batch_size=64)


class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.linear = Linear(196608,10)

    def forward(self,input):
        output = self.linear(input)
        return output


mymodule = MyModule()
for data in dataloader:
    imgs,targets = data
    print(imgs.shape)
    #output = torch.reshape(imgs,(1,1,1,-1))
    output = torch.flatten(imgs)#与上面的reshape功能一样
    print(output.shape)
    output1 = mymodule(output)
    print(output1.shape)